import os
import json

from tqdm import tqdm
from typing import Any, Dict, List
from .i_reader import IReader
from .structure import RtkData


class RtkReader(IReader[RtkData]):
    def __init__(self, cfg: Dict[str, Any]):
        self.parallel = cfg.get('parallel')
        self.input_tensor_output_folder = cfg.get('input_tensor_output_folder')

    def _get_file_list(self, folder_path: str) -> List[str]:
        return [
            os.path.join(folder_path, entry.name)
            for entry in os.scandir(folder_path)
            if entry.name.endswith('.rtk')
        ]

    def _read_single_file(self, file_path: str) -> RtkData:
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
            data['heading'] = -data['heading']
            data['height'] = data['height'] / 1000.0
            return RtkData(**data)

    def process_all_files(self, folder_path, frame_ids: List[str]) -> Dict[str, RtkData]:
        all_files = self._get_file_list(folder_path)
        
        files = []
    
        for frame_id in frame_ids:
            full_path = os.path.join(folder_path, f"{frame_id}.rtk")
            
            if full_path in all_files:
                files.append(full_path)
        
        if self.parallel:
            from multiprocessing import Pool, cpu_count

            with Pool(processes=max(1, cpu_count() - 1)) as pool:
                rtk_data_list = list(tqdm(
                    pool.imap(self._read_single_file, files),
                    total=len(files),
                    desc="Processing RTK files"
                ))
                
                results = {
                    os.path.splitext(os.path.basename(files[i]))[0]: rtk_data_list[i] 
                    for i in range(len(files))
                }
        else:
            results = {
                os.path.splitext(os.path.basename(f))[0]: self._read_single_file(f) 
                for f in tqdm(files, desc="Reading rtk files")
            }
        
        return results

    def read_tensor_rtk(self, folder_path, batch_name) -> Dict[str, RtkData]:
        pt_files = [
            f for f in os.listdir(self.input_tensor_output_folder)
            if f.endswith('.pt') and f.startswith(batch_name)
        ]
        
        rtk_files_to_read = []
        
        for pt_file in pt_files:
            parts = pt_file.split('_', 2)
            if len(parts) == 3:
                rtk_filename = f"{parts[2].replace('.pt', '')}.rtk"
                rtk_files_to_read.append(rtk_filename)
        
        files = [
            os.path.join(folder_path, f)
            for f in rtk_files_to_read
            if os.path.exists(os.path.join(folder_path, f))
        ]
        
        if self.parallel:
            from multiprocessing import Pool, cpu_count

            with Pool(processes=max(1, cpu_count() - 1)) as pool:
                rtk_data_list = list(tqdm(
                    pool.imap(self._read_single_file, files),
                    total=len(files),
                    desc="Processing tensor RTK files"
                ))
                
                results = {
                    os.path.splitext(os.path.basename(files[i]))[0]: rtk_data_list[i] 
                    for i in range(len(files))
                }
        else:
            results = {
                os.path.splitext(os.path.basename(f))[0]: self._read_single_file(f) 
                for f in tqdm(files, desc="Reading tensor rtk files")
            }
        
        return results